import torch
import torch.nn as nn


class SeModule(nn.Module):
    def __init__(self, in_channel, reduction=4):
        super(SeModule, self).__init__()
        expand_size = max(in_channel // reduction, 8)
        self.se = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channel, expand_size, kernel_size=1, bias=False),
            nn.BatchNorm2d(expand_size),
            nn.ReLU(inplace=True),
            nn.Conv2d(expand_size, in_channel, kernel_size=1, bias=False),
            nn.Hardsigmoid()
        )

    def forward(self, x):
        return x * self.se(x)



if __name__ == '__main__':
    for name, p in nn.LSTM(4,40).named_parameters():
        print(name, p.size())
    x = torch.randn((2,3,32,32))
    models = SeModule(3)
    print(models(x).size())